import torch

"""
因果掩码
"""


def generate_square_subsequent_mask(size: int) -> torch.Tensor:
    """
    首先创建一个size×size的矩阵，所有元素初始化为负无穷大；
    然后通过torch.triu()函数保留上三角部分（对角线偏移1），其余位置设为0；
    最终返回该掩码矩阵，用于Transformer中的序列解码时防止未来信息泄露。
    """
    mask = torch.full((size, size), float('-inf'))
    """
    triu (upper triangle) 函数返回输入张量的上三角部分，其余部分置为0
    """
    mask = torch.triu(mask, diagonal=1)
    return mask


if __name__ == '__main__':
    print(generate_square_subsequent_mask(6))
